// "Copyright [2021] <Copyright Shengkai Lin>"

#ifndef TASK_SCHEDULER_TASK_GENERATOR_H_
#define TASK_SCHEDULER_TASK_GENERATOR_H_

#include <vector>
#include "ai_sim.h"
#include <unordered_map>
#include "network_controller/topo_manager.h"
#include "resource_scheduler/resource_scheduler.h"
#include "task_scheduler/task_type/allreduce.h"
#include "network_controller/network_controller.h"


namespace ai_sim {

class TaskGenerator : public omnetpp::cSimpleModule{
 private:
       TaskType tmp_tasktype;
       int maxnumGpus;
       int maxnumBatches;
       int modelsize;
       double compute_time;
       double task_interval;
       int total_tasks;
       TopoManager *topos;
       ResourceScheduler *resource_schuler;
       NetworkController *net_control;
       std::unordered_map<int, taskboard> taskboards;


 public:
        void initialize() override;
        void handleMessage(omnetpp::cMessage *msg) override;
        void finish() override;
        
        Task gen_fully_utilized_task();
            
        void gen_comm_pair(taskboard&, AllreduceAlgorithm*);
        void gen_flows(taskboard&, AllreduceAlgorithm*);
        void allocAllFlows(taskboard&);
        void release_task(TaskMsg*);
        void try_release_task(const Task&, TaskMsg*);
        void get_statistics(const taskboard&);
};  // namespace ai_sim
}
#endif  // TASK_SCHEDULER_TASK_GENERATOR_H_
